--[[
  1. key : use this key to find the token bucket in redis
  2. intervalPerPermit : time interval in millis between two token permits,equals interval / limit
  3. refillTime : current timestamp (millsecond)
  4. burstTokens : initial value
  5. limit : capacity
  6. interval : limit time interval
]] --

local key, intervalPerPermit, refillTime, burstTokens = KEYS[1], tonumber(ARGV[1]), tonumber(ARGV[2]), tonumber(ARGV[3])
local limit, interval = tonumber(ARGV[4]), tonumber(ARGV[5])
local bucket = redis.call('hgetall', key)

local currentTokens

if table.maxn(bucket) == 0 then
    -- first check if bucket not exists, if yes, create a new one with full burstTokens, then grant access
    currentTokens = burstTokens
    redis.call('hset', key, 'lastRefillTime', refillTime)
    -- set expire time when init the key
    -- redis.call('pexpire', key, interval * 1.5)
elseif table.maxn(bucket) >= 4 then
    -- if bucket exists, first we try to refill the token bucket

    local lastRefillTime, tokensRemaining = tonumber(bucket[2]), tonumber(bucket[4])

    if refillTime > lastRefillTime then
        -- if refillTime larger than lastRefillTime, we should refill the token buckets
        -- calculate the interval between refillTime and lastRefillTime
        -- if the result is bigger than the interval of the token bucket,
        -- refill the tokens to capacity limit;
        -- else calculate how much tokens should be refilled
        local intervalSinceLast = refillTime - lastRefillTime
        if intervalSinceLast > interval then
            currentTokens = burstTokens
            redis.call('hset', key, 'lastRefillTime', refillTime)
        else
            local grantedTokens = math.floor(intervalSinceLast / intervalPerPermit)
            if grantedTokens > 0 then
                -- ajust lastRefillTime, we want shift left the refill time.
                local padMillis = math.fmod(intervalSinceLast, intervalPerPermit)
                redis.call('hset', key, 'lastRefillTime', refillTime - padMillis)
            end
            currentTokens = math.min(grantedTokens + tokensRemaining, limit)
        end
    else
        -- if not, it means some other operation later than this call made the call first.
        -- there is no need to refill the tokens.
        currentTokens = tokensRemaining
    end
end

assert(currentTokens >= 0)
redis.call('pexpire', key, interval * 1.5)
if currentTokens == 0 then
    redis.call('hset', key, 'tokensRemaining', currentTokens)
    return 0
else
    redis.call('hset', key, 'tokensRemaining', currentTokens - 1)
    return currentTokens
end

